import cv2
import numpy as np
from typing import List, Tuple
from sklearn.cluster import DBSCAN

def cluster_contours(contours, eps=50, min_samples=2):
    """
    使用DBSCAN聚类算法将相近的轮廓聚合
    
    Args:
        contours: 轮廓列表
        eps: DBSCAN中邻域的大小
        min_samples: 形成聚类的最小样本数
    
    Returns:
        clustered_contours: 聚类后的轮廓列表
    """
    if len(contours) == 0:
        return []
    
    # 计算每个轮廓的质心
    centroids = []
    for contour in contours:
        M = cv2.moments(contour)
        if M["m00"] != 0:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
        else:
            # 如果轮廓面积为0，使用边界框中心
            x, y, w, h = cv2.boundingRect(contour)
            cx, cy = x + w//2, y + h//2
        centroids.append([cx, cy])
    
    centroids = np.array(centroids)
    
    # 使用DBSCAN聚类
    clustering = DBSCAN(eps=eps, min_samples=min_samples).fit(centroids)
    labels = clustering.labels_
    
    # 合并同一聚类中的轮廓
    clustered_contours = []
    unique_labels = set(labels)
    
    for label in unique_labels:
        if label == -1:  # 噪声点，单独处理
            noise_indices = np.where(labels == label)[0]
            for idx in noise_indices:
                clustered_contours.append(contours[idx])
        else:
            # 合并同一聚类的轮廓
            cluster_indices = np.where(labels == label)[0]
            merged_contour = merge_contours([contours[i] for i in cluster_indices])
            if merged_contour is not None:
                clustered_contours.append(merged_contour)
    
    return clustered_contours

def merge_contours(contours_to_merge):
    """
    合并多个轮廓为一个轮廓
    """
    if len(contours_to_merge) == 0:
        return None
    if len(contours_to_merge) == 1:
        return contours_to_merge[0]
    
    # 将所有轮廓点合并
    all_points = []
    for contour in contours_to_merge:
        all_points.extend(contour.reshape(-1, 2))
    
    all_points = np.array(all_points)
    
    # 计算凸包作为合并后的轮廓
    hull = cv2.convexHull(all_points)
    
    return hull

def process_app_screenshot(image_path: str):
    """
    处理手机APP截图，检测大白块区域
    """
    # 1. 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError("无法读取图像文件")
    
    height, width = img.shape[:2]
    screen_area = height * width
    min_area_threshold = screen_area / 8  # 面积阈值
    
    # 2. 灰度化
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 3. 边缘检测 (多种方案可选)
    edges = detect_edges(gray)
    
    # 4. 查找轮廓
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 4.5. 轮廓聚类合并
    print(f"原始轮廓数量: {len(contours)}")
    clustered_contours = cluster_contours(contours, eps=100, min_samples=2)
    print(f"聚类后轮廓数量: {len(clustered_contours)}")
    
    # 5. 筛选大部件并判断是否为白块
    large_white_blocks = []
    
    # 创建用于绘制轮廓的图像副本
    contour_img = img.copy()
    
    for contour in clustered_contours:
        area = cv2.contourArea(contour)
        x, y, w, h = cv2.boundingRect(contour)
        # 绘制轮廓
        cv2.drawContours(contour_img, [contour], -1, (0, 255, 0), 3)
        # 绘制边界框
        cv2.rectangle(contour_img, (x, y), (x+w, y+h), (255, 0, 0), 2)
        # 添加标签
        cv2.putText(contour_img, f"White Block {len(large_white_blocks)}", 
                    (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
        if area > min_area_threshold:
            # 判断是否为白块
            if is_white_block(img, x, y, w, h):
                large_white_blocks.append({
                    'bbox': (x, y, w, h),
                    'area': area,
                    'contour': contour
                })
                
    
    # 保存绘制了轮廓的图像
    import os
    from datetime import datetime
    
    # 获取输出目录
    output_dir = "/Users/atria/Desktop/coding/hypiumKea/Kea2/loadertests"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 生成时间戳文件名
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    contour_filename = f"{output_dir}/white_blocks_contours_{timestamp}.jpg"
    
    # 保存图像
    cv2.imwrite(contour_filename, contour_img)
    print(f"轮廓图像已保存到: {contour_filename}")
    
    return large_white_blocks, edges, contour_img

def detect_edges(gray_img):
    """
    边缘检测 - 提供多种算法选择
    """
    # 方案1: Canny边缘检测 (推荐)
    # 先高斯模糊降噪
    blurred = cv2.GaussianBlur(gray_img, (5, 5), 0)
    edges = cv2.Canny(blurred, 50, 150)
    
    # 方案2: Sobel边缘检测
    # sobel_x = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=3)
    # sobel_y = cv2.Sobel(gray_img, cv2.CV_64F, 0, 1, ksize=3)
    # edges = np.sqrt(sobel_x**2 + sobel_y**2).astype(np.uint8)
    
    # 方案3: Laplacian边缘检测
    # edges = cv2.Laplacian(gray_img, cv2.CV_64F, ksize=3)
    # edges = np.absolute(edges).astype(np.uint8)
    
    return edges

def is_white_block(img, x, y, w, h, white_threshold=200):
    """
    判断指定区域是否为白块
    """
    # 提取感兴趣区域(ROI)
    roi = img[y:y+h, x:x+w]
    
    # 方法1: 基于平均亮度判断
    gray_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    mean_brightness = np.mean(gray_roi)
    
    # 方法2: 基于白色像素占比
    white_pixels = np.sum(gray_roi > white_threshold)
    total_pixels = gray_roi.size
    white_ratio = white_pixels / total_pixels
    
    # 方法3: 基于颜色方差 (白块颜色较均匀)
    color_variance = np.var(gray_roi)
    
    # 综合判断条件
    is_white = (
        mean_brightness > white_threshold and  # 平均亮度高
        white_ratio > 0.8 and                # 白色像素占比高
        color_variance < 500                  # 颜色变化小(均匀)
    )
    
    return is_white

def visualize_results(img, large_white_blocks, edges):
    """
    可视化检测结果
    """
    # 绘制检测到的白块
    result_img = img.copy()
    for block in large_white_blocks:
        x, y, w, h = block['bbox']
        cv2.rectangle(result_img, (x, y), (x+w, y+h), (0, 255, 0), 2)
        cv2.putText(result_img, f"White Block", (x, y-10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
    
    return result_img

# 使用示例
if __name__ == "__main__":
    image_path = "/Users/atria/Desktop/coding/hypiumKea/Kea2/loadertests/20250806-230942.jpeg"  # 你的截图路径
    
    try:
        white_blocks, edges, contour_img = process_app_screenshot(image_path)
        
        print(f"检测到 {len(white_blocks)} 个大白块:")
        for i, block in enumerate(white_blocks):
            x, y, w, h = block['bbox']
            print(f"白块 {i+1}: 位置({x}, {y}), 大小({w}x{h}), 面积={block['area']}")
        
        # 可视化结果
        original_img = cv2.imread(image_path)
        result_img = visualize_results(original_img, white_blocks, edges)
        
        # 显示结果
        cv2.imshow('Original', original_img)
        cv2.imshow('Edges', edges)
        cv2.imshow('Detection Result', result_img)
        cv2.imshow('Contours with White Blocks', contour_img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        
    except Exception as e:
        print(f"处理错误: {e}")